import cv2
import numpy as np

# 初始化摄像头
cap = cv2.VideoCapture(0)

# 轨迹点存储列表
trajectory = []
max_trajectory_length = 64  # 保留最近的64个轨迹点

# 定义红色HSV范围
lower_red = np.array([0, 120, 70])
upper_red = np.array([10, 255, 255])
lower_red2 = np.array([170, 120, 70])
upper_red2 = np.array([180, 255, 255])

while True:
    ret, frame = cap.read()
    if not ret:
        break
    
    # 转换为HSV颜色空间
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    # 创建红色掩膜
    mask1 = cv2.inRange(hsv, lower_red, upper_red)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask = cv2.bitwise_or(mask1, mask2)
    
    # 形态学操作（降噪）
    kernel = np.ones((5,5), np.uint8)
    mask = cv2.erode(mask, kernel, iterations=2)
    mask = cv2.dilate(mask, kernel, iterations=2)
    
    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if contours:
        # 找到最大轮廓
        max_contour = max(contours, key=cv2.contourArea)
        ((x, y), radius) = cv2.minEnclosingCircle(max_contour)
        
        if radius > 10:  # 过滤小的噪点
            # 绘制圆心和轮廓
            center = (int(x), int(y))
            cv2.circle(frame, center, int(radius), (0, 255, 0), 2)
            cv2.circle(frame, center, 3, (0, 0, 255), -1)
            
            # 添加轨迹点
            trajectory.append(center)
            if len(trajectory) > max_trajectory_length:
                trajectory.pop(0)
    
    # 绘制轨迹
    for i in range(1, len(trajectory)):
        if trajectory[i-1] is None or trajectory[i] is None:
            continue
        thickness = int(np.sqrt(max_trajectory_length / (i+1)) * 2.5)
        cv2.line(frame, trajectory[i-1], trajectory[i], (255, 0, 0), thickness)
    
    cv2.imshow('Tracking', frame)
    
    # 按q退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()